'''
Created on Jan 16 2012

@author: oabalbin

This script takes
input: a list of snps positions by sample
output a file like CHR  POS    SP1   SP2   SP3
                    1   345    A/T   A/T   A/C
'''
import os
import sys
import subprocess
import numpy as np
from optparse import OptionParser
from collections import deque, defaultdict

                
def read_positions(ifile):
    '''
    '''
    ifile = open(ifile)
    for l in ifile:
        if l.startswith('#'):
            continue
        fields=l.strip('\n').split(',')
        fields.pop()
    
    ifile.close()
    return fields

def convert_thsnvs(thsnvs):
    '''
    chrM@2354|C>T|1/1,
    '''
    al1 = set()
    al2 = set()
    for snv in thsnvs:
        s = snv.split('|')
        p,a,b = s[0],s[1].split('>')[0],s[1].split('>')[1]
        al1.add( p+'|'+a+'>*' )
        al2.add( p+'|*>'+b )
    
    return al1,al2
    

def get_thsnvs_pos(thsnvs):
    '''
    chrM@2354|C>T|1/1,
    '''
    al1 = set()
    for snv in thsnvs:
        s = snv.split('|')
        p = s[0]
        al1.add( p )
    return al1

def get_ibd_snvs(pos_file,thsnvs):
    '''
    chrM@2354|C>T|1/1,
    '''
    ifile = open(pos_file)
    for l in ifile:
        if l.startswith('#'):
            continue
        fields=l.strip('\n').split(',')

    ibd_snvs=set()
    pos = get_thsnvs_pos(thsnvs)
    print "Getting the postions", len(pos)
    for i in fields:
        p0=i.split('|')[0]
        if p0 in pos:
            ibd_snvs.add( i )
    return ibd_snvs
    
    
def isec_arrays(filea,fileb):
    '''
    '''
    m1=set(read_positions(filea))
    m2=set(read_positions(fileb))
    #print sys.getsizeof(m1), sys.getsizeof(m2)
    isec=m1.intersection(m2)        
    return isec

def isec_snvs_annot_union(snvs_union, snvs_annot_dbs):
    '''
    This functions obtains all query snvs that are 
    common to a particular database, ex: dbSNP
    '''
    pm = defaultdict()
    for db, snv_db in snvs_annot_dbs.iteritems():
        pm[db] = set(snvs_union).intersection(set(snv_db))
           
    return pm

def union_snvs_dict(union_snvs):
    '''
    It gives an index to each snp. 
    So it can be used later to find arrays
    '''
    pm = defaultdict()
    for i, snv in enumerate(union_snvs):
        pm[snv]=i
    return pm


def read_patient_files(file,datadir):
    ifile=open(file)
    pat=defaultdict()
    
    for l in ifile:
        f=l.strip('\n').split('\t')
        pat[f[0]]=os.path.join(datadir,f[2])
    ifile.close()
    return pat


def ibd_children(patient_files,datadir):
    '''
    '''
    children = ['broth1','broth2','case']
    First=True
    for c in children:
        print os.path.join(datadir,patient_files[c])
        thsnvs = read_positions(os.path.join(datadir,patient_files[c]))
        # return 2 sets
        al1, al2 = convert_thsnvs(thsnvs)
        if First:
            al1_snvs_isec,al2_snvs_isec= al1,al2
            First=False
        else:
            al1_snvs_isec.intersection_update( al1 )
            al2_snvs_isec.intersection_update( al2 )
    return al1_snvs_isec,al2_snvs_isec


def children_snps(patient_files):
    '''
    '''
    children = ['broth1','broth2','case']
    First=True
    for c in children:
        thsnvs = get_thsnvs_pos(read_positions(patient_files[c]))
        if First:
            s=thsnvs
            First=False
        else:
            s.union(thsnvs)
    return s



def get_positions_dict(thsnvs,prepos=None):
    '''
    chrM@2354|C>T|1/1,
    return dictionary of positions -> genotype.
    '''
    snps=defaultdict()
    for snv in thsnvs:
        s = snv.split('|')
        p,g = s[0],s[2]
        if prepos is not None and p in prepos:
            snps[p]=g.split('/')[0]+g.split('/')[1]
        else:
            snps[p]=g.split('/')[0]+g.split('/')[1]
    return snps


class child:
    def __init__(self,snps,hsnps):
        self.mysnps=snps # my snps positions
        self.hsnps=snps.intersection(hsnps) # snps positions shared with both parents
        self.prsnps=snps.difference(hsnps) #private snps positions
        print "My number of snps positions ", len(self.mysnps)
        print "The number of positions shared with my parents ", len(self.hsnps)
        print "The number of private postions is ", len(self.prsnps)
        
class family:
    def __init__(self,snps_files):
        #self.father=read_positions(snps_files['father'])
        #self.mother=read_positions(snps_files['mother'])
        #self.scase=read_positions(snps_files['case'])
        #self.sbroth1=read_positions(snps_files['broth1'])
        #self.sbroth2=read_positions(snps_files['broth2'])
        self.snps_files=snps_files
        
        self.have_family()
        
        # create following attributes:
        #self.pasnps=set()
        #self.common_family_snps=set()
        #self.case= child(get_thsnvs_pos(self.scase), self.pasnps)
        #self.broth1= child(get_thsnvs_pos(self.sbroth1), self.pasnps)
        #self.broth2= child(get_thsnvs_pos(self.sbroth2), self.pasnps)

        
        
    def parental_snps(self):
        father = get_thsnvs_pos(read_positions(self.snps_files['father']))
        mother = get_thsnvs_pos(read_positions(self.snps_files['mother']))
        self.pasnps = father.union(mother)
    
    def transmitted_snps(self):
        # intersection snps accross all children and parents
        c = self.children_snps()
        self.common_family_snps =  self.pasnps.intersection(c)
        
    def children_snps(self):
        '''
        returns common snps to all children
        '''
        children = ['broth1','broth2','case']
        First=True
        for c in children:
            thsnvs = get_thsnvs_pos(read_positions(self.snps_files[c]))
            if First:
                s=thsnvs
                First=False
            else:
                s.intersection_update(thsnvs)
        return s
    
    def have_family(self):
        self.parental_snps()
        self.transmitted_snps()
        print "The number of parental snps is ", len(self.pasnps)
        print "The number of transmitted snps is ", len(self.common_family_snps)
        #self.case= child(get_thsnvs_pos(self.scase), self.pasnps)
        #self.broth1= child(get_thsnvs_pos(self.sbroth1), self.pasnps)
        #self.broth2= child(get_thsnvs_pos(self.sbroth2), self.pasnps)
        
    def assume_homozygous(self,ref):
        # if the snps is not present, assume that it was homozygous for the reference
        return ref[0]+ref[0]
    
    def get_inheritance_state(self,ofile):
        '''
        '''
        ofile=open(ofile,'w')
        names = ['broth1','broth2','case']
        children=[]
        for c in names:
            children.append( get_positions_dict(read_positions(self.snps_files[c]), self.common_family_snps) )
        
        #case_common_snps = get_positions_dict(self.scase,self.common_family_snps)
        #bro1_common_snps = get_positions_dict(self.sbroth1,self.common_family_snps)
        #bro2_common_snps = get_positions_dict(self.sbroth2,self.common_family_snps)
        
        #case_common_snps = get_positions_dict(read_positions(self.snps_files['case']),self.common_family_snps)
        mother_snps = get_positions_dict(read_positions(self.snps_files['mother']),self.common_family_snps)
        father_snps = get_positions_dict(read_positions(self.snps_files['father']),self.common_family_snps)
        
        for snp in self.common_family_snps:
            allg=[]
            ref='00' # children[0][snp]
            for child in children:
                allg.append( child[snp] )
            GA='NA'
            try:
                fa=father_snps[snp]
            except KeyError:
                print "[SNP NOT IN FATHER] homozygous assumed ",snp
                fa=self.assume_homozygous(ref)
                GA='FA'
            try:
                mo=mother_snps[snp]
            except KeyError:
                print "[SNP NOT IN MOTHER] homozygous assumed ",snp
                mo=self.assume_homozygous(ref)
                GA='MA'
            
            state=self.snp_inheratence_state(fa,mo,allg)
            chr,start,end,info,strand=snp.split('@')[0],int(snp.split('@')[1])-1,int(snp.split('@')[1])+1,\
                                      state+'|'+fa+mo+'|'+",".join(allg).replace(',',';')+'|'+GA,'+'
            ol=[chr,start,end,info,strand]
            print ol
            ofile.write(",".join(map(str,ol)).replace(',','\t')+'\n')

    def get_all_inheritance_state(self,ofile):
        '''
        '''
        ofile=open(ofile,'w')
        names = ['broth1','broth2','case']
        children=[]
        for c in names:
            children.append( get_positions_dict(read_positions(self.snps_files[c])) )
        
        #case_common_snps = get_positions_dict(self.scase,self.common_family_snps)
        #bro1_common_snps = get_positions_dict(self.sbroth1,self.common_family_snps)
        #bro2_common_snps = get_positions_dict(self.sbroth2,self.common_family_snps)
        
        #case_common_snps = get_positions_dict(read_positions(self.snps_files['case']),self.common_family_snps)
        mother_snps = get_positions_dict(read_positions(self.snps_files['mother']))
        father_snps = get_positions_dict(read_positions(self.snps_files['father']))
        
        all_snps=[]
        for snp in self.pasnps:
            allg=[]
            ref='00' # children[0][snp]
            CA,nca='',0
            for child in children:
                try:
                    gc = child[snp]
                    CA +='NA'
                except KeyError:
                    #print "[SNP NOT IN CHILD] homozygous assumed ",snp
                    gc=self.assume_homozygous(ref)
                    CA+='CA'
                    nca+=1
                allg.append( gc )
            if nca == len(children):
                #print "[SNP NOT IN CHILDREN] Skip Inheritance State ",snp
                continue
            
            GA='NA'
            try:
                fa=father_snps[snp]
            except KeyError:
                #print "[SNP NOT IN FATHER] homozygous assumed ",snp
                fa=self.assume_homozygous(ref)
                GA='FA'
            try:
                mo=mother_snps[snp]
            except KeyError:
                #print "[SNP NOT IN MOTHER] homozygous assumed ",snp
                mo=self.assume_homozygous(ref)
                GA='MA'
            
            state=self.snp_inheratence_state(fa,mo,allg)
            #print fa, mo, allg,state
            callInfo=str(state)+'|'+fa+mo+'|'+",".join(allg).replace(',',';')+'|'+CA+'|'+GA
            chr,start,end,info,strand=snp.split('@')[0],int(snp.split('@')[1])-1,int(snp.split('@')[1])+1,\
                                      callInfo,'+'
            il=[chr,start,end,info,strand]
            ofile.write(",".join(map(str,il)).replace(',','\t')+'\n')
            
            ol = chr.split('chr')[1]
            if ol=='X':
                ol=23
            elif ol=='Y':
                ol=24
            elif ol=='M': # CHRM M
                continue

            all_snps.append((int(ol),int(start),int(end),info,strand))
        
        return all_snps
    
    
    def snp_inheratence_state(self,fa,mo,ch):
        '''
        Inheritance state: I=identical, HP=haploidentical Paternal,
        HM=Haploidentical Maternal, A=Ambiguos (HP/HM), P=Private
        '''
        ms=3       
        if fa == '00' and mo == '01':
            aa=ch.count('00')
            ab=ch.count('01')
            if aa==ms or ab==ms:
                return 'I'
            else:
                return 'HP'
        elif fa == '01' and mo == '00':
            aa=ch.count('00')
            ab=ch.count('01')
            #ba=ch.count('01')
            bb=ch.count('11')
            if aa==ms or ab==ms:
                return 'I'
            elif ab > 0:
                return 'HM'
            elif bb >0:
                return 'E'
            
        elif fa == '01' and mo == '01':
            aa=ch.count('00')
            bb=ch.count('11')
            ab=ch.count('01')
            if aa==ms or bb==ms:
                return 'I'
            else:
                return 'A' # It is an ambiguos state it could be HP/HM
        elif fa == '11' and mo == '01':
            bb=ch.count('11')
            ab=ch.count('01')
            ba=ch.count('10')
            if bb==ms:
                return 'I'
            elif ba > 0 or ab > 0:
                return 'HP'
            '''
            else:
                return 'A' # specially when the snp was not called in all the children
            '''
        elif fa == '01' and mo == '11':
            bb=ch.count('11')
            ab=ch.count('01')
            ba=ch.count('10')
            if bb==ms:
                return 'I'
            elif ab > 0 or ba > 0:
                return 'HM'
            '''
            else:
                return 'A' # specially when the snp was not called in all the children
            '''

        elif fa == '00' and mo == '11':
            ab=ch.count('01')
            if ab==ms:
                return 'I'
            else:
                return 'E' # It is an ambiguos state it could be Error
            
        elif fa == '11' and mo == '00':
            ab=ch.count('01')
            if ab==ms:
                return 'I'
            else:
                return 'E' # It is an ambiguos state it could be Error
        
        elif fa=='11' and mo =='11':
            bb = ch.count('11')
            if bb==ms:
                return 'I'
            else:
                return 'E'
        elif fa == '00' and mo == '00':
            aa = ch.count('00')
            if aa==ms:
                return 'I'
            else:
                return 'E'
            
        
        

        
        


if __name__ == '__main__':
    '''
    Write to an array the positions in a particualr vcf.
    
    '''
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-f", "--pos_files", action="append", type="str",
                            dest="pos_files",
                            help="others_vcf_files. Pass vcf files without .gz extention even if the file .gz is there")
    
    optionparser.add_option("-p", "--patient_files", type="str", dest="patient_files",
                            help="others_vcf_files. Pass vcf files without .gz extention even if the file .gz is there")


    optionparser.add_option("-o", "--out_file", type="str",dest="ofile",
                            help="others_vcf_files. Pass vcf files without .gz extention even if the file .gz is there")
    
    optionparser.add_option("-b", "--base", dest="base", action="store_true", default=False,
                            help="determine if the ouput file is position or pos.base")
    optionparser.add_option("-g", "--geno", dest="geno", action="store_true", default=False,
                            help="determine if the ouput file is position or pos.base.geno")

    
    
    (options, args) = optionparser.parse_args()
    
    
    '''
    father    51    FMPN51.gatk_snps.b.filtered.model.pos.base.geno
    mother    65    FMPN65.gatk_snps.b.filtered.model.pos.base.geno
    broth1    52    FMPN52.gatk_snps.b.filtered.model.pos.base.geno
    case    63    FMPN63.gatk_snps.b.filtered.model.pos.base.geno
    broth2    53    FMPN53.gatk_snps.b.filtered.model.pos.base.geno
        
    '''
    datadir='/exds/users/oabalbin/projects/fmpn/data/gatk_calls'    
    ofile=os.path.join(datadir,'all_snps_inheratence_state.bed')
    patients_files = read_patient_files(options.patient_files,datadir)
    Gang = family(patients_files)
    all_snvs = Gang.get_all_inheritance_state(ofile)
    #Gang.get_inheritance_state(ofile)
    
    
    # For the snvs shared across all individuals, isec set
    nrows=len(all_snvs)
    ncols=1
    #mytype=[('CHRM',int),('POS',int),('IBD',int),('MUT','S3'),('GENO','S3'),('SAMPLE',int)]
    #mytype_ibd=[('CHRM',int),('POS',int),('IBD',int)]
    #HM|0100|01;01;00|NANACA|MA
    mytype=[('CHRM',int),('START',int),('END',int),('INFO','S27'),('STRAND','S2')]
  
    #location_arr=np.empty(( nrows, ncols),dtype=mytype)
    '''
    for snv in isec_set:
        loc,base,geno = snv.split('|')[0],snv.split('|')[1].replace('>','/'),snv.split('|')[2]
        chr,pos=loc.split('@')[0],loc.split('@')[1]
        chrN=chr.replace('chr','')
        if chrN=='X':
            chrN=23
        elif chrN=='Y':
            chrN=24
        elif chrN=='M':
            chrN=25
        start, end =int(pos)-1,int(pos)+1
        info = base+'|'+geno
        all_snvs.append((chrN,start,end))#,info)) # This indixing is important, 0 means is shared for all samples
        all_ibd.append((chrN,start,end))#,info))
        #print ol, snv
        #print location_arr[i,:]
        i+=1
      '''
    #sys.exit(0)
    # Sort the snv table
    location_arr=np.array( all_snvs,dtype=mytype )
    location_arr_sorted = np.sort(location_arr, order=['CHRM', 'START'])
    
    ########### Writing everything out
    of2=open(options.ofile+'_all','w')
       
    chrs = np.array(range(1,26))
    i=1
    pad=1e4
    markers=0
    of2=open(options.ofile+'_hap_all.txt','w')
    
    for i in chrs:
        sub_array=location_arr_sorted[location_arr_sorted['CHRM']==i]
        
        print sub_array
        for j in range(len(sub_array)):
            ol=list(sub_array[j])
            
            if ol[0]==23:
                ol[0]='X'
            elif ol[0]==24:
                ol[0]='Y'
            elif ol[0]==25: # CHRM M
                continue
            #print ",".join(map(str,ol)).replace(',','\t')
            if j==0:
                start=min(sub_array['START'])
                state=[]
            else:
                if ol[1] <= start + pad:
                    markers+=1
                    state.append(ol[3].split('|')[0])
                    continue
                else: 
                    c = list(sub_array[j-1])
                    a = (start + pad) - c[1]
                    b = ol[1] - (start + pad)
                    if a < b:
                        end = c[1]
                    else:
                        end = ol[1]
                    
                    hd=state.count('HP')
                    hm=state.count('HM')
                    id=state.count('I')
                    e=state.count('E')
                    a=state.count('A')
                    n=state.count('None')
                    info='%d|HP%d;HM%d;I%d;E%d;A%d;N%d'%(markers,hd,hm,id,e,a,n)
                    l = [ol[0],start,end,info,0,'+']
                    of2.write('chr'+",".join(map(str,l)).replace(',','\t')+'\n')
                    start = ol[1]
                    markers=0
                    state=[]
    of2.close()

